

package keeno.usap.analysis.pta.plugin.reflection;

import keeno.usap.analysis.pta.core.cs.element.CSObj;
import keeno.usap.analysis.pta.core.heap.Descriptor;
import keeno.usap.analysis.pta.core.heap.HeapModel;
import keeno.usap.analysis.pta.core.heap.Obj;
import keeno.usap.analysis.pta.core.solver.Solver;
import keeno.usap.analysis.pta.plugin.util.CSObjs;
import keeno.usap.ir.exp.ClassLiteral;
import keeno.usap.ir.stmt.Invoke;
import keeno.usap.language.annotation.Annotation;
import keeno.usap.language.classes.ClassMember;
import keeno.usap.language.classes.ClassNames;
import keeno.usap.language.classes.JClass;
import keeno.usap.language.classes.JField;
import keeno.usap.language.classes.JMethod;
import keeno.usap.language.type.ClassType;
import keeno.usap.language.type.Type;
import keeno.usap.language.type.TypeSystem;

import javax.annotation.Nullable;
import java.util.Objects;
import java.util.Set;

/**
 * Helps manage reflection and annotation objects.
 */
class MetaObjHelper {

    /**
     * Descriptor for the meta objects generated by reflection inference.
     */
    private static final Descriptor REFLECTION_DESC = () -> "ReflectionMetaObj";

    /**
     * Descriptor for the meta objects generated by reflection log.
     */
    private static final Descriptor REFLECTION_LOG_DESC = () -> "ReflectionLogMetaObj";

    /**
     * Descriptor for the array objects generated by
     * Class.get[Declared][Constructor|Method|Field]s().
     */
    private static final Descriptor REFLECTION_ARRAY_DESC = () -> "ReflectionMetaObjArray";

    /**
     * Descriptor for unknown meta objects.
     */
    private static final Descriptor UNKNOWN_REFLECTION_DESC = () -> "UnknownReflectionMetaObj";

    /**
     * Descriptor for annotation objects.
     */
    private static final Descriptor ANNOTATION_DESC = () -> "AnnotationObj";

    private final TypeSystem typeSystem;

    private final HeapModel heapModel;

    final ClassType clazz;

    final ClassType constructor;

    final ClassType method;

    final ClassType field;

    MetaObjHelper(Solver solver) {
        typeSystem = solver.getTypeSystem();
        heapModel = solver.getHeapModel();
        clazz = typeSystem.getClassType(ClassNames.CLASS);
        constructor = typeSystem.getClassType(ClassNames.CONSTRUCTOR);
        method = typeSystem.getClassType(ClassNames.METHOD);
        field = typeSystem.getClassType(ClassNames.FIELD);
    }

    /**
     * Given a JClass, Type, or ClassMember, return the corresponding meta object.
     * @throws IllegalArgumentException if type of {@code classOrTypeOrMember}
     * is neither {@link JClass}, {@link Type}, nor {@link ClassMember}.
     */
    Obj getMetaObj(Object classOrTypeOrMember) {
        return getMetaObj(classOrTypeOrMember, REFLECTION_DESC);
    }

    Obj getLogMetaObj(Object classOrTypeOrMember) {
        return getMetaObj(classOrTypeOrMember, REFLECTION_LOG_DESC);
    }

    boolean isLogMetaObj(CSObj obj) {
        return CSObjs.hasDescriptor(obj, REFLECTION_LOG_DESC);
    }

    private Obj getMetaObj(Object classOrTypeOrMember, Descriptor desc) {
        if (classOrTypeOrMember instanceof JClass jclass) {
            ClassLiteral classLiteral = ClassLiteral.get(jclass.getType());
            return desc.equals(REFLECTION_LOG_DESC)
                    ? heapModel.getMockObj(desc, classLiteral, clazz)
                    : heapModel.getConstantObj(classLiteral);
        } else if (classOrTypeOrMember instanceof Type type) {
            ClassLiteral classLiteral = ClassLiteral.get(type);
            return desc.equals(REFLECTION_LOG_DESC)
                    ? heapModel.getMockObj(desc, classLiteral, clazz)
                    : heapModel.getConstantObj(classLiteral);
        } else if (classOrTypeOrMember instanceof JMethod m) {
            if (m.isConstructor()) {
                return heapModel.getMockObj(desc, classOrTypeOrMember, constructor);
            } else {
                return heapModel.getMockObj(desc, classOrTypeOrMember, method);
            }
        } else if (classOrTypeOrMember instanceof JField) {
            return heapModel.getMockObj(desc, classOrTypeOrMember, field);
        } else {
            throw new IllegalArgumentException(
                    "Expected JClass or ClassMember," + " given " + classOrTypeOrMember);
        }
    }

    Obj getMetaObjArray(Invoke invoke) {
        assert Set.of("getConstructors", "getDeclaredConstructors",
                "getMethods", "getDeclaredMethods",
                "getFields", "getDeclaredFields").contains(invoke.getMethodRef().getName());
        return heapModel.getMockObj(REFLECTION_ARRAY_DESC, invoke,
                invoke.getMethodRef().getReturnType(), invoke.getContainer());
    }

    /**
     * @return an unknown class generated at {@code invoke}.
     */
    Obj getUnknownClass(Invoke invoke) {
        return heapModel.getMockObj(UNKNOWN_REFLECTION_DESC, invoke, clazz, false);
    }

    Obj getUnknownMethod(Invoke invoke, @Nullable JClass clazz, @Nullable String name) {
        MethodInfo methodInfo = new MethodInfo(invoke, clazz, name);
        return heapModel.getMockObj(UNKNOWN_REFLECTION_DESC, methodInfo, method,
                invoke.getContainer(), false);
    }

    MethodInfo getMethodInfo(CSObj csObj) {
        return (MethodInfo) csObj.getObject().getAllocation();
    }

    boolean isUnknownMetaObj(CSObj csObj) {
        return CSObjs.hasDescriptor(csObj, UNKNOWN_REFLECTION_DESC);
    }

    Obj getAnnotationObj(Annotation annotation) {
        ClassType type = typeSystem.getClassType(annotation.getType());
        Objects.requireNonNull(type,
                "Annotation type" + annotation.getType() + " is not found");
        return heapModel.getMockObj(ANNOTATION_DESC, annotation, type);
    }
}
